package com.jq.jqaicode.common.langgraphDemo;

import dev.langchain4j.community.model.dashscope.QwenStreamingChatModel;
import dev.langchain4j.data.message.AiMessage;
import dev.langchain4j.data.message.ChatMessage;
import dev.langchain4j.data.message.UserMessage;
import dev.langchain4j.model.chat.request.ChatRequest;
import dev.langchain4j.model.chat.request.ChatRequestParameters;
import lombok.extern.slf4j.Slf4j;
import org.bsc.langgraph4j.GraphStateException;
import org.bsc.langgraph4j.action.EdgeAction;
import org.bsc.langgraph4j.action.NodeAction;
import org.bsc.langgraph4j.langchain4j.generators.StreamingChatGenerator;
import org.bsc.langgraph4j.langchain4j.serializer.std.LC4jStateSerializer;
import org.bsc.langgraph4j.langchain4j.tool.LC4jToolService;
import org.bsc.langgraph4j.prebuilt.MessagesState;
import org.bsc.langgraph4j.prebuilt.MessagesStateGraph;
import org.bsc.langgraph4j.streaming.StreamingOutput;

import static org.bsc.langgraph4j.action.AsyncEdgeAction.edge_async;
import static org.bsc.langgraph4j.StateGraph.START;
import static org.bsc.langgraph4j.StateGraph.END;
import static org.bsc.langgraph4j.action.AsyncNodeAction.node_async;

import java.util.Map;

@Slf4j
public class StreamingGraphAppDemo {


    public static void main(String[] args) throws GraphStateException {
//        var generator = StreamingChatGenerator.<AgentState>builder()
//                .mapResult( r -> Map.of( "content", r.aiMessage().text() ) )
//                .build();

        var stateSerializer = new LC4jStateSerializer<MessagesState<ChatMessage>>( MessagesState::new );


        var model = QwenStreamingChatModel.builder()
                .apiKey("sk-a56c6cac116e4db8a8492f18aebcab88")
                .modelName("qwen-max")
                .build();

//        var request = ChatRequest.builder()
//                .messages( UserMessage.from("Tell me a joke") )
//                .build();
//
//        model.chat(request, generator.handler() );

            // setup tools
        var tools = LC4jToolService.builder()
                .toolsFromObject( new SearchTool() )
                .build();

        NodeAction<MessagesState<ChatMessage>> callModel = state -> {
           // log.info( "CallModel" );

            var generator = StreamingChatGenerator.<MessagesState<ChatMessage>>builder()
                    .mapResult( response -> Map.of("messages", response.aiMessage()) )
                    .startingNode("agent")
                    .startingState(state)
                    .build();

            var parameters = ChatRequestParameters.builder()
                    .toolSpecifications(tools.toolSpecifications())
                    .build();

            var request = ChatRequest.builder()
                    .parameters(parameters)
                    .messages( state.messages() )
                    .build();

            model.chat( request, generator.handler() );

            return Map.of("_streaming_messages", generator);
        };

        // Route Message
        EdgeAction<MessagesState<ChatMessage>> routeMessage = state -> {

            ChatMessage lastMessage = state.lastMessage()
                    .orElseThrow(() -> (new IllegalStateException("last message not found!")));


            //log.info("routeMessage:\n{}", lastMessage );
            if (lastMessage instanceof AiMessage message) {
                // If tools should be called
                if (message.hasToolExecutionRequests()) {
                    return "next";
                }
            }

            // If no tools are called, we can finish (respond to the user)
            return "exit";
        };


//        for( var r : generator ) {
//            //log.info( "{}", r);
//            System.out.println(r);
//        }


        // Invoke Tool
        NodeAction<MessagesState<ChatMessage>> invokeTool = state -> {

            var lastMessage = state.lastMessage()
                    .orElseThrow(() -> (new IllegalStateException("last message not found!")));

            //log.info("invokeTool:\n{}", lastMessage );

            if (lastMessage instanceof AiMessage lastAiMessage) {

                var result = tools.execute(lastAiMessage.toolExecutionRequests(), null)
                        .orElseThrow(() -> (new IllegalStateException("no tool found!")));

                return Map.of("messages", result);

            }

            throw new IllegalStateException("invalid last message");
        };

        // Define Graph
        var workflow = new MessagesStateGraph<ChatMessage>(stateSerializer)
                .addNode("agent", node_async(callModel))
                .addNode("tools", node_async(invokeTool))
                .addEdge(START, "agent")
                .addConditionalEdges("agent",
                        edge_async(routeMessage),
                        Map.of("next", "tools", "exit", END))
                .addEdge("tools", "agent");

        var app = workflow.compile();

        for( var out : app.stream( Map.of( "messages", UserMessage.from( "获取上海的天气")) ) ) {
            if( out instanceof StreamingOutput streaming ) {
                log.info( "StreamingOutput{node={}, chunk={} }", streaming.node(), streaming.chunk() );
            }
            else {
                log.info( "{}", out );
            }
        }
    }
}
